from pathlib import Path
import yaml
import argparse
from typing import List, Tuple

from asr.frontend.default import DefaultFrontend
from asr.specaug.specaug import SpecAug
from layers.global_mvn import GlobalMVN
from asr.encoder.conformer_encoder import ConformerEncoder
from asr.decoder.transformer_decoder import TransformerDecoder
from asr.ctc import CTC
from asr.espnet_model import ESPnetASRModel

class ASRConformerModel():

    def build_model(self, args):
        if isinstance(args.token_list, str):
            with open(args.token_list, encoding="utf-8") as f:
                token_list = [line.rstrip() for line in f]
            args.token_list = list(token_list)
        elif isinstance(args.token_list, (tuple, list)):
            token_list = list(args.token_list)
        else:
            raise RuntimeError("token_list must be str or list")
        
        # 1. frontend
        if args.input_size is None:
            frontend = DefaultFrontend
            input_size = frontend.output_size()
        else:
            pass
        # 2. data augmentation for spectrogram
        if args.specaug is not None:
            specaug = SpecAug
        else:
            specaug = None
        # 3. Normalization layer
        if args.normalize is not None:
            normalize = GlobalMVN
        else:
            normalize = None
        # 4. pre-encoder input block
        # None
        # 4. Encoder
        encoder = ConformerEncoder
        # 5. post-encoder block
        # 5. Decoder
        if args.decoder == "transformer":
            # TODO:add input augments
            decoder = TransformerDecoder()
        else:
            pass
        # 6. CTC
        # TODO:add input augments
        ctc = CTC()
        # 7. build model
        # TODO:add input augments
        model = ESPnetASRModel()



    def build_model_from_file(self, config_file, model_file):

        if config_file is None:
            assert model_file is not None, "'model_file' must be provided."
            config_file = Path(model_file).parent / "config.yaml"
        else:
            config_file = Path(config_file)
        
        with config_file.open("r", encoding="utf-8") as f:
            args = yaml.safe_load(f)
        args = argparse.Namespace(**args)